#include "sbox.h"
# include<time.h>

// 字节转化为字
#ifndef GM_GET_UINT32_BE
#define GM_GET_UINT32_BE(n, b, i)                       \
{                                                       \
    (n) = ( (WORD) (b)[(i)    ] << 24 )             \
        | ( (WORD) (b)[(i) + 1] << 16 )             \
        | ( (WORD) (b)[(i) + 2] <<  8 )             \
        | ( (WORD) (b)[(i) + 3]       );            \
}
#endif

// 字转化为字节
#ifndef GM_PUT_UINT32_BE
#define GM_PUT_UINT32_BE(n, b ,i)                       \
{                                                       \
    (b)[(i)    ] = (unsigned char) ( (n) >> 24 );       \
    (b)[(i) + 1] = (unsigned char) ( (n) >> 16 );       \
    (b)[(i) + 2] = (unsigned char) ( (n) >>  8 );       \
    (b)[(i) + 3] = (unsigned char) ( (n)       );       \
}
#endif

#define  GM_SHL(x,n) (((x) & 0xFFFFFFFF) << (n % 32))
#define GM_ROTL(x,n) (GM_SHL((x),n) | ((x) >> (32 - (n % 32))))

#define  GM_SHL256(x,n) (((x) & 0xFFFFFFFFFFFFFFFF) << (n % 64))
#define GM_ROTL256(x,n) (GM_SHL256((x),n) | ((x) >> (64 - (n % 64))))

typedef unsigned long long int WORD;
typedef WORD KEYWORD;

#define PLAINTEXTBLOCKSIZE 512

typedef struct {
    unsigned char key[32];
    KEYWORD rk[32];
    KEYWORD dec_key[32];
}SM4_ECB;

// S盒 仅用于密钥拓展使用
unsigned char GM_SM4_SBOX[256] = {
	0xd6, 0x90, 0xe9, 0xfe, 0xcc, 0xe1, 0x3d, 0xb7, 0x16, 0xb6, 0x14, 0xc2, 0x28, 0xfb, 0x2c, 0x05,
	0x2b, 0x67, 0x9a, 0x76, 0x2a, 0xbe, 0x04, 0xc3, 0xaa, 0x44, 0x13, 0x26, 0x49, 0x86, 0x06, 0x99,
	0x9c, 0x42, 0x50, 0xf4, 0x91, 0xef, 0x98, 0x7a, 0x33, 0x54, 0x0b, 0x43, 0xed, 0xcf, 0xac, 0x62,
	0xe4, 0xb3, 0x1c, 0xa9, 0xc9, 0x08, 0xe8, 0x95, 0x80, 0xdf, 0x94, 0xfa, 0x75, 0x8f, 0x3f, 0xa6,
	0x47, 0x07, 0xa7, 0xfc, 0xf3, 0x73, 0x17, 0xba, 0x83, 0x59, 0x3c, 0x19, 0xe6, 0x85, 0x4f, 0xa8,
	0x68, 0x6b, 0x81, 0xb2, 0x71, 0x64, 0xda, 0x8b, 0xf8, 0xeb, 0x0f, 0x4b, 0x70, 0x56, 0x9d, 0x35,
	0x1e, 0x24, 0x0e, 0x5e, 0x63, 0x58, 0xd1, 0xa2, 0x25, 0x22, 0x7c, 0x3b, 0x01, 0x21, 0x78, 0x87,
	0xd4, 0x00, 0x46, 0x57, 0x9f, 0xd3, 0x27, 0x52, 0x4c, 0x36, 0x02, 0xe7, 0xa0, 0xc4, 0xc8, 0x9e,
	0xea, 0xbf, 0x8a, 0xd2, 0x40, 0xc7, 0x38, 0xb5, 0xa3, 0xf7, 0xf2, 0xce, 0xf9, 0x61, 0x15, 0xa1,
	0xe0, 0xae, 0x5d, 0xa4, 0x9b, 0x34, 0x1a, 0x55, 0xad, 0x93, 0x32, 0x30, 0xf5, 0x8c, 0xb1, 0xe3,
	0x1d, 0xf6, 0xe2, 0x2e, 0x82, 0x66, 0xca, 0x60, 0xc0, 0x29, 0x23, 0xab, 0x0d, 0x53, 0x4e, 0x6f,
	0xd5, 0xdb, 0x37, 0x45, 0xde, 0xfd, 0x8e, 0x2f, 0x03, 0xff, 0x6a, 0x72, 0x6d, 0x6c, 0x5b, 0x51,
	0x8d, 0x1b, 0xaf, 0x92, 0xbb, 0xdd, 0xbc, 0x7f, 0x11, 0xd9, 0x5c, 0x41, 0x1f, 0x10, 0x5a, 0xd8,
	0x0a, 0xc1, 0x31, 0x88, 0xa5, 0xcd, 0x7b, 0xbd, 0x2d, 0x74, 0xd0, 0x12, 0xb8, 0xe5, 0xb4, 0xb0,
	0x89, 0x69, 0x97, 0x4a, 0x0c, 0x96, 0x77, 0x7e, 0x65, 0xb9, 0xf1, 0x09, 0xc5, 0x6e, 0xc6, 0x84,
	0x18, 0xf0, 0x7d, 0xec, 0x3a, 0xdc, 0x4d, 0x20, 0x79, 0xee, 0x5f, 0x3e, 0xd7, 0xcb, 0x39, 0x48
};

#define GM_SM4_NON_LINEAR_OP(in, out) \
{                                                           \
	(out)  = ( GM_SM4_SBOX[ ((in) >> 24) & 0x0FF ] ) << 24; \
	(out) |= ( GM_SM4_SBOX[ ((in) >> 16) & 0x0FF ] ) << 16; \
	(out) |= ( GM_SM4_SBOX[ ((in) >> 8)  & 0x0FF ] ) << 8;  \
	(out) |= ( GM_SM4_SBOX[ (in)         & 0x0FF ] );       \
}

/* System parameter */
unsigned int GM_SM4_FK[4] = {0xA3B1BAC6, 0x56AA3350, 0x677D9197, 0xB27022DC};

/* fixed parameter */
static const unsigned long GM_SM4_CK[32] =
{
	0x00070e15, 0x1c232a31, 0x383f464d, 0x545b6269,
	0x70777e85, 0x8c939aa1, 0xa8afb6bd, 0xc4cbd2d9,
	0xe0e7eef5, 0xfc030a11, 0x181f262d, 0x343b4249,
	0x50575e65, 0x6c737a81, 0x888f969d, 0xa4abb2b9,
	0xc0c7ced5, 0xdce3eaf1, 0xf8ff060d, 0x141b2229,
	0x30373e45, 0x4c535a61, 0x686f767d, 0x848b9299,
	0xa0a7aeb5, 0xbcc3cad1, 0xd8dfe6ed, 0xf4fb0209,
	0x10171e25, 0x2c333a41, 0x484f565d, 0x646b7279
};

static void SM4KeySchedule(const unsigned char * key, unsigned int * rk)
{
	unsigned tmp1, tmp2, K[36];
	int i;

	// transfer input
	// K0, K1, K2, K3
    GM_GET_UINT32_BE( K[0], key, 0);
    GM_GET_UINT32_BE( K[1], key, 4);
    GM_GET_UINT32_BE( K[2], key, 8);
    GM_GET_UINT32_BE( K[3], key, 12);

    K[0] ^= GM_SM4_FK[0];
    K[1] ^= GM_SM4_FK[1];
    K[2] ^= GM_SM4_FK[2];
    K[3] ^= GM_SM4_FK[3];

	// rki = Ki+4
	for(i = 0; i < 32; i++)
	{
		// non linear's input
		tmp1 = K[i + 1] ^ K[i + 2] ^ K[i + 3] ^ GM_SM4_CK[i];

		// non linear operation
		GM_SM4_NON_LINEAR_OP(tmp1, tmp2);

		// linear operation
		rk[i] = K[i + 4] = K[i] ^ (tmp2 ^ GM_ROTL(tmp2, 13) ^ GM_ROTL(tmp2, 23));
	}
}

static void KeyExtension(unsigned char * key, KEYWORD * rk) {
    unsigned int rk1[32];
    unsigned int rk2[32];

    SM4KeySchedule(key, rk1);
    SM4KeySchedule(key + 16, rk2);

    for(int i = 0;i < 32; i++) {
        rk[i] = (KEYWORD)rk1[i] + (KEYWORD)rk2[i];
    }
}

static void KeyRevser(KEYWORD * rk, KEYWORD * dec_rk) {
    for (int i = 32 - 1, j = 0; i >= 0 ; i --, j++) {
        dec_rk[j] = rk[i];
    }
}

typedef struct {
    WORD X1;
    WORD X2;
    WORD X3;
    WORD X4;
}BLOCK;

void Reverse(BLOCK * target) {
    WORD tmp;
    for(int i = 0; i < PLAINTEXTBLOCKSIZE/sizeof(BLOCK); i++ ) {
        tmp = target[i].X1;
        target[i].X1 = target[i].X4;
        target[i].X4 = tmp;

        tmp = target[i].X2;
        target[i].X2 = target[i].X3;
        target[i].X3 = tmp;
    }
}


static void F256(unsigned char *plaintext, KEYWORD * rk) {
    // 每次处理256字节的数据
    BLOCK * block = (BLOCK *)plaintext;
    WORD buf[256 / sizeof(WORD)] = {0};
    
    // 32轮函数
    for(int r = 0; r < 32; r ++) {
        for(int i = 0; i < PLAINTEXTBLOCKSIZE / sizeof(BLOCK); i++) {
            buf[i] = block[i].X2 ^ block[i].X3 ^ block[i].X4 ^ rk[r];
        }

        //search sbox
        long_bits in = BUILDLONGBITS();
        long_bits out = BUILDLONGBITS();

        Block2LongBits((unsigned char *)buf, &in);
        LongBitsSbox(in, &out);
        LongBits2Block(out, (unsigned char *)buf);

        //ROL and xor
        for(int i = 0; i < PLAINTEXTBLOCKSIZE / sizeof(BLOCK); i++){
            buf[i] = block[i].X1 ^ (GM_ROTL256(buf[i], 2)) ^ (GM_ROTL256(buf[i], 10)) ^ (GM_ROTL256(buf[i], 18)) ^ (GM_ROTL256(buf[i], 24));
        }

        // change the plaintext
        memcpy(plaintext, plaintext + sizeof(WORD), PLAINTEXTBLOCKSIZE - sizeof(WORD));

        for(int i = 0; i < PLAINTEXTBLOCKSIZE / sizeof(BLOCK); i++) {
            block[i].X4 = buf[i];
        }
        memset(buf, 0, 256);
    }

    // revser
    Reverse(block);

}


unsigned char * EncryptECB(SM4_ECB * ctx, unsigned char * plaintext, int len) {
    KEYWORD *rk = ctx->rk;
    for(int i = 0; i < len / PLAINTEXTBLOCKSIZE; i++) F256(plaintext + PLAINTEXTBLOCKSIZE*i, rk);
    return plaintext;
}

unsigned char * DecryptECB(SM4_ECB * ctx, unsigned char * cipher, int len) {
    KEYWORD *dec_rk = ctx->dec_key;
    for(int i = 0; i < len / PLAINTEXTBLOCKSIZE; i++) F256(cipher + PLAINTEXTBLOCKSIZE*i, dec_rk); 
    return cipher;
}

void BUILDSM4_ECB(SM4_ECB * sm4_ctx, unsigned char * key) { 
    memcpy(sm4_ctx->rk, key, 32);
    KeyExtension(key, sm4_ctx->rk);
    KeyRevser(sm4_ctx->rk, sm4_ctx->dec_key);
}


int test256() {
    //64:32k   *32 : 16k    16:8K
    unsigned char plaintext[PLAINTEXTBLOCKSIZE];
    unsigned char * tmp = "SM4 is a symmetric block cipher algorithm that operates on 128-bit blocks and uses a secret key. It follows a 32-round Feistel network structure, employing a combination of substitution, permutation, and XOR operations,  as well as round function iterations. SM4 provides strong cryptographic security, ensuring confidentiality, integrity, and authenticity. With widespread adoption and support, SM4 is utilized in various applications such as secure communication, protection, and privacy-preserving protocols.";
    for(unsigned int i=0;i<1;i++){
        memcpy(plaintext+i*PLAINTEXTBLOCKSIZE, tmp, PLAINTEXTBLOCKSIZE);
    }
    unsigned char key[] = {0x21, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef, 0xfe, 0xdc, 0xba, 0x98, 0x76, 0x54, 0x32, 0x10, 0x21, 0x23, 0x45, 0x88, 0x89, 0xab, 0xcd, 0xef, 0xde, 0xdc, 0xba, 0x98, 0x76, 0x54, 0x32, 0x41};
    printf("\n================================================================================================================================================================= \n");
    SM4_ECB ctx;
    clock_t begin, end;
    BUILDSM4_ECB(&ctx, key);
    
    printf("\nmessage: \n  ");
    for (int i = 0; i < PLAINTEXTBLOCKSIZE; i++) printf("%c", plaintext[i]);
    printf("\n");
    printf("\nkey: \n  ");
    for(int i = 0; i < 64; i++) {
        printf("%2x ", key[i]);
        if(i && (i+1) % 32 == 0){
            printf("\n  ");
        }
    }
    begin = clock();
    EncryptECB(&ctx, plaintext, PLAINTEXTBLOCKSIZE);
    end = clock();
    double t1 = (double)(end-begin)*1000/CLOCKS_PER_SEC;
    
    printf("\nCiphertext: \n  ");
    for (int i = 0; i < PLAINTEXTBLOCKSIZE; i++) {
        printf("%2x", plaintext[i]);
        if((i+1) % 64 == 0){
            printf("\n  ");
        }
    }
    printf("\n");
    begin = clock();
    DecryptECB(&ctx, plaintext, PLAINTEXTBLOCKSIZE);
    end = clock();
    double t2 = (double)(end-begin)*1000/CLOCKS_PER_SEC;
    printf("plaintext: \n  ");
    for (int i = 0; i < PLAINTEXTBLOCKSIZE; i++) printf("%c", plaintext[i]);
    printf("\n\n");
    printf("encryption time: %f ms\n", t1);
    printf("decryption time: %f ms \n", t2);
    printf("\n================================================================================================================================================================= \n");
    // puts("");

}

int main() {
    test256();
    return 0;
}


